import json
from pathlib import Path
from collections import Counter
from collections import OrderedDict
from argparse import ArgumentParser
from common.configs.path import paths
from common.configs.tools import logger, init_logger


title_counter = Counter()
content_counter = Counter()
tc_counter = Counter()  # title + content


def read_json(path):
    with open(path) as f:
        for line in f:
            parse_json(line)


def parse_json(line):
    text = json.loads(line.strip())
    title_list = text['title'].split(' ')
    content_list = text['content'].split(' ')

    global title_counter
    global content_counter
    global tc_counter

    title_counter += Counter(title_list)
    content_counter += Counter(content_list)
    tc_counter += Counter(title_list + content_list)


def save_dict(path, dt):
    with open(path, 'w') as f:
        json.dump(dt, f)
    print(path, 'saved.')


def main():
    parser = ArgumentParser(description='')
    parser.add_argument('--file', type=Path, help='file path',
                        default=paths['pre_train_data'])
    parser.add_argument(
        '--obj', type=str, help='title (t) or content (c) or title & content (tc)')
    parser.add_argument('--top', type=int,
                        help='output top N char', default=10)

    args = parser.parse_args()

    topn = args.top

    logger = init_logger(log_file=paths['log_dir'])

    read_json(args.file)

    global title_counter
    global content_counter
    global tc_counter

    logger.info('title: {}'.format(title_counter.most_common(topn)))
    logger.info('content: {}'.format(content_counter.most_common(topn)))
    logger.info('title & content: {}'.format(tc_counter.most_common(topn)))
    logger.info('total word: {}'.format(len(tc_counter)))

    save_dict(paths['output'] / 'title_count.json', dict(title_counter))
    save_dict(paths['output'] / 'content_count.json', dict(content_counter))
    save_dict(paths['output'] / 'tc_count.json', dict(tc_counter))

    # pri


if __name__ == '__main__':
    main()
